import torch
import torch.nn as nn


class SmallBasicBlock(nn.Module):
    def __init__(self, ch_in, ch_out):
        super(SmallBasicBlock, self).__init__()
        self.block = nn.Sequential(
            nn.Conv2d(ch_in, ch_out // 4, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(ch_out // 4, ch_out // 4, kernel_size=(3, 1), padding=(1, 0)),
            nn.ReLU(),
            nn.Conv2d(ch_out // 4, ch_out // 4, kernel_size=(1, 3), padding=(0, 1)),
            nn.ReLU(),
            nn.Conv2d(ch_out // 4, ch_out, kernel_size=1),
        )

    def forward(self, x):
        return self.block(x)


class OldLPRNet(nn.Module):
    def __init__(self, class_num=66, dropout_rate=0.5):
        super().__init__()
        self.backbone = nn.Sequential(  # [bs,3,24,94]
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1),  # 0  -> [bs,64,22,92]
            nn.BatchNorm2d(num_features=64),  # 1  -> [bs,64,22,92]
            nn.ReLU(),  # 2  -> [bs,64,22,92]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 1, 1)),  # 3  -> [bs,64,20,90]
            SmallBasicBlock(ch_in=64, ch_out=128),  # 4  -> [bs,128,20,90]
            nn.BatchNorm2d(num_features=128),  # 5  -> [bs,128,20,90]
            nn.ReLU(),  # 6  -> [bs,128,20,90]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(2, 1, 2)),  # 7  -> [bs,64,18,44]
            SmallBasicBlock(ch_in=64, ch_out=256),  # 8  -> [bs,256,18,44]
            nn.BatchNorm2d(num_features=256),  # 9  -> [bs,256,18,44]
            nn.ReLU(),  # 10 -> [bs,256,18,44]
            SmallBasicBlock(ch_in=256, ch_out=256),  # 11 -> [bs,256,18,44]
            nn.BatchNorm2d(num_features=256),  # 12 -> [bs,256,18,44]
            nn.ReLU(),  # 13 -> [bs,256,18,44]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(4, 1, 2)),  # 14 -> [bs,64,16,21]
            nn.Dropout(dropout_rate),  # 0.5 dropout rate                          # 15 -> [bs,64,16,21]
            nn.Conv2d(in_channels=64, out_channels=256, kernel_size=(1, 4), stride=1),  # 16 -> [bs,256,16,18]
            nn.BatchNorm2d(num_features=256),  # 17 -> [bs,256,16,18]
            nn.ReLU(),  # 18 -> [bs,256,16,18]
            nn.Dropout(dropout_rate),  # 0.5 dropout rate                                  19 -> [bs,256,16,18]
            nn.Conv2d(in_channels=256, out_channels=class_num, kernel_size=(13, 1), stride=1),
            # class_num=66  20  -> [bs,66,4,18]
            nn.BatchNorm2d(num_features=class_num),  # 21 -> [bs,66,4,18]
            nn.ReLU(),  # 22 -> [bs,66,4,18]
        )
        self.container = nn.Sequential(
            nn.Conv2d(in_channels=448 + class_num, out_channels=class_num, kernel_size=(1, 1), stride=(1, 1)),
            # nn.BatchNorm2d(num_features=self.class_num),
            # nn.ReLU(),
            # nn.Conv2d(in_channels=self.class_num, out_channels=self.lpr_max_len+1, kernel_size=3, stride=2),
            # nn.ReLU(),
        )

    def forward(self, x):
        keep_features = list()
        for i, layer in enumerate(self.backbone.children()):
            x = layer(x)
            if i in [2, 6, 13, 22]:  # 2: [bs,64,22,92]  6:[bs,128,20,90] 13:[bs,256,18,44] 22:[bs,68,4,18]
                keep_features.append(x)
        global_context = list()  # keep_features: [bs,64,22,92]  [bs,128,20,90] [bs,256,18,44] [bs,68,4,18]
        for i, f in enumerate(keep_features):
            if i in [0, 1]:  # [bs,64,22,92] -> [bs,64,4,18]  [bs,128,20,90] -> [bs,128,4,18]
                f = nn.AvgPool2d(kernel_size=5, stride=5)(f)
            if i in [2]:  # [bs,256,18,44] -> [bs,256,4,18]
                f = nn.AvgPool2d(kernel_size=(4, 10), stride=(4, 2))(f)
            global_context.append(torch.div(f, torch.mean(torch.pow(f, 2))))  # 所有元素求平方求平均以这个均值
        # [bs,64,4,18]+[bs,128,4,18]+[bs,256,4,18]+[bs,68,4,18]=[bs,516,4,18] -> [bs, 66, 4, 18]  -> [bs, 66, 18]
        return torch.mean(self.container(torch.cat(global_context, 1)), dim=2)


class MidLPRNet(nn.Module):
    def __init__(self, class_num=66, dropout_rate=0.5):
        super().__init__()
        self.backbone = nn.Sequential(  # [bs,3,24,94]
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1),  # 0  -> [bs,64,22,92]
            nn.BatchNorm2d(num_features=64),  # 1  -> [bs,64,22,92]
            nn.ReLU(),  # 2  -> [bs,64,22,92]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 1, 1)),  # 3  -> [bs,64,20,90]
            SmallBasicBlock(ch_in=64, ch_out=128),  # 4  -> [bs,128,20,90]
            nn.BatchNorm2d(num_features=128),  # 5  -> [bs,128,20,90]
            nn.ReLU(),  # 6  -> [bs,128,20,90]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(2, 1, 2)),  # 7  -> [bs,64,18,44]
            SmallBasicBlock(ch_in=64, ch_out=256),  # 8  -> [bs,256,18,44]
            nn.BatchNorm2d(num_features=256),  # 9  -> [bs,256,18,44]
            nn.ReLU(),  # 10 -> [bs,256,18,44]
            SmallBasicBlock(ch_in=256, ch_out=256),  # 11 -> [bs,256,18,44]
            nn.BatchNorm2d(num_features=256),  # 12 -> [bs,256,18,44]
            nn.ReLU(),  # 13 -> [bs,256,18,44]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(4, 1, 2)),  # 14 -> [bs,64,16,21]
            # nn.Dropout(dropout_rate),  # 0.5 dropout rate                          # 15 -> [bs,64,16,21]
            nn.BatchNorm2d(num_features=64),
            nn.Conv2d(in_channels=64, out_channels=256, kernel_size=(1, 4), stride=1),  # 16 -> [bs,256,16,18]
            nn.BatchNorm2d(num_features=256),  # 17 -> [bs,256,16,18]
            nn.ReLU(),  # 18 -> [bs,256,16,18]
            # nn.Dropout(dropout_rate),  # 0.5 dropout rate                                  19 -> [bs,256,16,18]
            nn.BatchNorm2d(num_features=256),
            nn.Conv2d(in_channels=256, out_channels=class_num, kernel_size=(13, 1), stride=1),
            # class_num=66  20  -> [bs,66,4,18]
            nn.BatchNorm2d(num_features=class_num),  # 21 -> [bs,66,4,18]
            nn.ReLU(),  # 22 -> [bs,66,4,18]
        )
        self.container = nn.Sequential(
            nn.Conv2d(in_channels=448 + class_num, out_channels=class_num, kernel_size=(1, 1), stride=(1, 1)),
            # nn.BatchNorm2d(num_features=self.class_num),
            # nn.ReLU(),
            # nn.Conv2d(in_channels=self.class_num, out_channels=self.lpr_max_len+1, kernel_size=3, stride=2),
            # nn.ReLU(),
        )

    def forward(self, x):
        keep_features = list()
        for i, layer in enumerate(self.backbone.children()):
            x = layer(x)
            if i in [2, 6, 13, 22]:  # 2: [bs,64,22,92]  6:[bs,128,20,90] 13:[bs,256,18,44] 22:[bs,68,4,18]
                keep_features.append(x)
        global_context = list()  # keep_features: [bs,64,22,92]  [bs,128,20,90] [bs,256,18,44] [bs,68,4,18]
        for i, f in enumerate(keep_features):
            if i in [0, 1]:  # [bs,64,22,92] -> [bs,64,4,18]  [bs,128,20,90] -> [bs,128,4,18]
                f = nn.AvgPool2d(kernel_size=5, stride=5)(f)
            if i in [2]:  # [bs,256,18,44] -> [bs,256,4,18]
                f = nn.AvgPool2d(kernel_size=(4, 10), stride=(4, 2))(f)
            global_context.append(torch.div(f, torch.mean(torch.pow(f, 2))))  # 所有元素求平方求平均以这个均值
        # [bs,64,4,18]+[bs,128,4,18]+[bs,256,4,18]+[bs,68,4,18]=[bs,516,4,18] -> [bs, 66, 4, 18]  -> [bs, 66, 18]
        return torch.mean(self.container(torch.cat(global_context, 1)), dim=2)


class NewLPRNet(nn.Module):
    def __init__(self, class_num=66):
        super().__init__()
        self.class_num = class_num
        self.backbone = nn.Sequential(  # [bs,3,24,94]
            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=1),  # 0  -> [bs,64,22,92]
            nn.BatchNorm2d(num_features=64),  # 1  -> [bs,64,22,92]
            nn.ReLU(),  # 2  -> [bs,64,22,92]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 1, 1)),  # 3  -> [bs,64,20,90]
            SmallBasicBlock(ch_in=64, ch_out=128),  # 4  -> [bs,128,20,90]
            nn.BatchNorm2d(num_features=128),  # 5  -> [bs,128,20,90]
            nn.ReLU(),  # 6  -> [bs,128,20,90]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(2, 1, 2)),  # 7  -> [bs,64,18,44]
            SmallBasicBlock(ch_in=64, ch_out=256),  # 8  -> [bs,256,18,44]
            nn.BatchNorm2d(num_features=256),  # 9  -> [bs,256,18,44]
            nn.ReLU(),  # 10 -> [bs,256,18,44]
            SmallBasicBlock(ch_in=256, ch_out=256),  # 11 -> [bs,256,18,44]
            nn.BatchNorm2d(num_features=256),  # 12 -> [bs,256,18,44]
            nn.ReLU(),  # 13 -> [bs,256,18,44]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(4, 1, 2)),  # 14 -> [bs,64,16,21]
            # nn.Dropout(dropout_rate),  # 0.5 dropout rate                          # 15 -> [bs,64,16,21]
            nn.BatchNorm2d(num_features=64),
            nn.Conv2d(in_channels=64, out_channels=256, kernel_size=(1, 4), stride=1),  # 16 -> [bs,256,16,18]
            nn.BatchNorm2d(num_features=256),  # 17 -> [bs,256,16,18]
            nn.ReLU(),  # 18 -> [bs,256,16,18]
            # nn.Dropout(dropout_rate),  # 0.5 dropout rate                                  19 -> [bs,256,16,18]
            nn.BatchNorm2d(num_features=256),
            nn.Conv2d(in_channels=256, out_channels=class_num, kernel_size=(13, 1), stride=1),
            # class_num=66  20  -> [bs,66,4,18]
            nn.BatchNorm2d(num_features=class_num),  # 21 -> [bs,66,4,18]
            nn.ReLU(),  # 22 -> [bs,66,4,18]
        )
        self.container = nn.Sequential(
            nn.Conv2d(in_channels=448 + class_num, out_channels=class_num, kernel_size=(1, 1), stride=(1, 1)),
            # nn.BatchNorm2d(num_features=self.class_num),
            # nn.ReLU(),
            # nn.Conv2d(in_channels=self.class_num, out_channels=self.lpr_max_len+1, kernel_size=3, stride=2),
            # nn.ReLU(),
        )

    def forward(self, x):
        keep_features = list()
        for i, layer in enumerate(self.backbone.children()):
            x = layer(x)
            if i in [2, 6, 13, 22]:  # 2: [bs,64,22,92]  6:[bs,128,20,90] 13:[bs,256,18,44] 22:[bs,68,4,18]
                keep_features.append(x)
        global_context = list()  # keep_features: [bs,64,22,92]  [bs,128,20,90] [bs,256,18,44] [bs,68,4,18]
        for i, f in enumerate(keep_features):
            if i in [0, 1]:  # [bs,64,22,92] -> [bs,64,4,18]  [bs,128,20,90] -> [bs,128,4,18]
                f = nn.AvgPool2d(kernel_size=5, stride=5)(f)
            if i in [2]:  # [bs,256,18,44] -> [bs,256,4,18]
                f = nn.AvgPool2d(kernel_size=(4, 10), stride=(4, 2))(f)
            global_context.append(torch.div(f, torch.mean(torch.pow(f, 2))))  # 所有元素求平方求平均以这个均值
        # [bs,64,4,18]+[bs,128,4,18]+[bs,256,4,18]+[bs,68,4,18]=[bs,516,4,18] -> [bs, 66, 4, 18]  -> [bs, 66, 18]
        return torch.mean(self.container(torch.cat(global_context, 1)), dim=2)


if __name__ == '__main__':
    from torch.autograd import Variable
    from thop import clever_format, profile
    from torch.utils.tensorboard import SummaryWriter


    def data(new_net=True):
        model = NewLPRNet() if new_net else OldLPRNet()
        dummy_input = torch.randn(1, 1, 24, 94) if new_net else torch.randn(1, 3, 24, 94)
        flops, params = profile(model, (dummy_input,), verbose=False)
        flops, params = clever_format([flops, params], "%.3f")
        print(f'Total GFLOPS:{flops}')
        print(f'Total params:{params}')
        # 测试网络，绘制网络图
        input_data = Variable(torch.rand(1, 1, 24, 94)) if new_net else Variable(torch.rand(1, 3, 24, 94))
        with SummaryWriter(comment='LPRNet') as w:
            w.add_graph(model, [input_data])


    data(new_net=False)
